from langchain_openai import ChatOpenAI
from langchain.prompts.chat import (
    ChatPromptTemplate,
    SystemMessagePromptTemplate,
    HumanMessagePromptTemplate,
)

from langchain import LLMChain
from langchain.schema import BaseOutputParser
import time

from config.model_config import get_chat_openai_xin, get_chat_openai_by_biaoshu, get_chat_openai_zhipu

# class CommaSeparatedListOutputParser(BaseOutputParser):
#     """将 LLM 调用的输出解析为逗号分隔的列表。"""
#     def parse(self, text: str):
#         """解析 LLM 调用的输出。"""
#         return text.strip().split(", ")
    

template = "您是一位有用的助手，可以生成逗号分隔的列表。用户将传入一个类别，您应该在该类别中生成 {number} 个对象，并以逗号分隔列表形式。仅返回逗号分隔的列表，仅此而已。"
system_message_prompt = SystemMessagePromptTemplate.from_template(template)
human_template = "{text}"
human_message_prompt = HumanMessagePromptTemplate.from_template(human_template)
print(human_message_prompt)


chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])
start_time = time.time()

chain = chat_prompt | get_chat_openai_zhipu(streaming=False)
res = chain.invoke({"number": 4, "text": "colors"})
print(res)
end_time = time.time()
elapsed_time = end_time - start_time

print(f"代码运行时间: {elapsed_time:.2f} 秒")